# from typing import override
from ..data.table_orm import get_kline_table_orm
from ..factors.base import FactorBase
import asyncio
from ccxt.pro import Exchange
from common.exchange import get_async_binance_portfolio_exchange_no_key
from common.variables import time_counts, TimeFrameType

# from pytypes import typechecked
import math
import pandas as pd


async def load_crypto_data(
    symbol: str,
    start_ms: int,
    end_ms: int,
    timeframe: TimeFrameType,
    exchange: Exchange = None,
) -> pd.DataFrame:
    # 这基本是单条线工作
    # data = pd.DataFrame(columns=["time", "open", "high", "low", "close", "volume"])
    data = []
    slide_time = start_ms
    exg = exchange
    if exchange is None:
        exg = get_async_binance_portfolio_exchange_no_key()
    while slide_time < end_ms:
        tmp_fetch = await exg.fetch_ohlcv(symbol, timeframe, slide_time, limit=5000)
        data.extend(tmp_fetch)
        if not tmp_fetch:
            break
        slide_time = tmp_fetch[-1][0] + 1
    if exchange is None:
        await exg.close()
    df = pd.DataFrame(data, columns=["time", "open", "high", "low", "close", "volume"])
    df.set_index("time", inplace=True)
    df = df[df.index < end_ms]
    return df


async def load_crypto_data_multitask(
    symbol: str,
    start_ms: int,
    end_ms: int,
    timeframe: TimeFrameType,
) -> pd.DataFrame:
    # 这里实现多条线工作
    exchange = get_async_binance_portfolio_exchange_no_key()
    maybe_counts = (
        int((end_ms - start_ms) / 1000 / time_counts[timeframe]) + 1
    )  # 估计的条数
    parts = min(10, math.ceil(maybe_counts / 1000))  # 保证每个任务不超过1000条
    delta = (end_ms - start_ms) // parts
    tasks = []
    for s in range(start_ms, end_ms, delta):
        tasks.append(
            load_crypto_data(symbol, s, min(s + delta, end_ms), timeframe, exchange)
        )
    results = await asyncio.gather(*tasks)
    await exchange.close()
    results = [x for x in results if not x.empty]
    if not results:
        return pd.DataFrame()
    df = pd.concat(results, axis=0)
    df = df.drop_duplicates().sort_index()
    df["symbol"] = symbol
    df.rename_axis("time", inplace=True)
    df.reset_index(inplace=True)
    return df


class FactorKline(FactorBase):
    bfname = "kline"
    fname_list = ["open", "high", "low", "close", "volume", "amount"]

    # @override
    def calculate(self, symbol: str, timeframe: TimeFrameType, start: int, end: int):
        results = asyncio.run(
            load_crypto_data_multitask(
                symbol,
                start,
                end,
                timeframe,
            )
        )
        orm = get_kline_table_orm(timeframe)
        results.rename(columns={"time": orm.xdatetime.name}, inplace=True)
        results[orm.symbol.name] = symbol
        results.set_index(orm.xdatetime.name, inplace=True)
        return results
